#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
YOLO标注工具 - 文件夹标注启动脚本
使用方法: python annotate_folder.py [文件夹名称]
"""

import os
import sys
from yolo_annotator_gui import run_yolo_annotator


def main():
    """主函数"""
    # 获取当前目录
    current_dir = os.getcwd()
    
    if len(sys.argv) > 1:
        # 如果提供了命令行参数，使用第一个参数作为文件夹名称
        folder_name = sys.argv[1]
        folder_path = os.path.join(current_dir, folder_name)
        
        if not os.path.exists(folder_path):
            print(f"错误: 文件夹 '{folder_name}' 不存在")
            print(f"当前目录: {current_dir}")
            print("请检查文件夹名称是否正确")
            return
            
        if not os.path.isdir(folder_path):
            print(f"错误: '{folder_name}' 不是一个文件夹")
            return
            
        # 检查文件夹中是否有图像文件
        image_files = [f for f in os.listdir(folder_path) 
                      if f.lower().endswith(('.jpg', '.jpeg', '.png', '.bmp', '.tiff'))]
        
        if not image_files:
            print(f"错误: 文件夹 '{folder_name}' 中没有找到图像文件")
            print("支持的图像格式: .jpg, .jpeg, .png, .bmp, .tiff")
            return
            
        print(f"找到 {len(image_files)} 张图像")
        print(f"启动标注工具，加载文件夹: {folder_path}")
        run_yolo_annotator(folder_path)
        
    else:
        # 没有提供参数，列出当前目录下的所有图像文件夹
        print("YOLO数据集标注工具")
        print("=" * 50)
        print(f"当前目录: {current_dir}")
        print()
        
        # 查找图像文件夹
        image_folders = []
        for item in os.listdir(current_dir):
            item_path = os.path.join(current_dir, item)
            if os.path.isdir(item_path):
                image_files = [f for f in os.listdir(item_path) 
                              if f.lower().endswith(('.jpg', '.jpeg', '.png', '.bmp', '.tiff'))]
                if image_files:
                    image_folders.append((item, item_path, len(image_files)))
        
        if not image_folders:
            print("未找到包含图像的文件夹")
            print("请确保当前目录下有包含图像文件的文件夹")
            print("支持的图像格式: .jpg, .jpeg, .png, .bmp, .tiff")
            print()
            print("使用方法:")
            print("  python annotate_folder.py <文件夹名称>")
            print("例如: python annotate_folder.py images")
            return
        
        # 显示找到的文件夹
        print("找到以下图像文件夹:")
        for i, (folder_name, folder_path, image_count) in enumerate(image_folders, 1):
            print(f"{i}. {folder_name} ({image_count}张图像)")
        
        print()
        print("使用方法:")
        print("  python annotate_folder.py <文件夹名称>")
        print("例如:")
        for folder_name, _, _ in image_folders:
            print(f"  python annotate_folder.py {folder_name}")


if __name__ == "__main__":
    main()
